# algorithms/gradient_boosting.py
import numpy as np
from collections import defaultdict
from .base_algorithm import BaseAlgorithm

try:
    import xgboost as xgb
    XGB_AVAILABLE = True
    print("[INFO] XGBoost可用")
except ImportError:
    print("[WARNING] XGBoost未安装，使用sklearn RandomForest替代")
    from sklearn.ensemble import RandomForestClassifier
    XGB_AVAILABLE = False

class GradientBoostingPredictor(BaseAlgorithm):
    """梯度提升算法 - 基于特征工程的机器学习方法"""
    
    def __init__(self, lottery_type='ssq'):
        super().__init__("Gradient-Boosting")
        self.lottery_type = lottery_type
        self.red_models = []
        self.blue_models = []
        self.feature_extractor = FeatureExtractor()
        
        # 初始化多个不同策略的模型
        self._initialize_models()
        
    def _initialize_models(self):
        """初始化多个不同配置的模型"""
        if XGB_AVAILABLE:
            # XGBoost模型配置
            configs = [
                {'n_estimators': 300, 'learning_rate': 0.1, 'max_depth': 6, 'subsample': 0.8},
                {'n_estimators': 400, 'learning_rate': 0.08, 'max_depth': 8, 'subsample': 0.85},
                {'n_estimators': 500, 'learning_rate': 0.12, 'max_depth': 5, 'subsample': 0.9}
            ]
            
            for i, config in enumerate(configs):
                model = xgb.XGBClassifier(
                    objective='multi:softprob',
                    random_state=42 + i,
                    verbosity=0,
                    **config
                )
                self.red_models.append(model)
        else:
            # RandomForest模型作为备选
            configs = [
                {'n_estimators': 200, 'max_depth': 10, 'min_samples_split': 5},
                {'n_estimators': 300, 'max_depth': 12, 'min_samples_split': 3},
                {'n_estimators': 250, 'max_depth': 8, 'min_samples_split': 4}
            ]
            
            for i, config in enumerate(configs):
                model = RandomForestClassifier(
                    random_state=42 + i,
                    n_jobs=-1,
                    **config
                )
                self.red_models.append(model)
    
    def train(self, historical_data):
        """训练梯度提升模型"""
        print(f"[INFO] 训练 {self.algorithm_name}...")
        
        try:
            # 特征工程
            features, labels = self.feature_extractor.extract_features_and_labels(historical_data)
            
            if len(features) == 0:
                print("[WARNING] 特征提取失败，使用随机初始化")
                self.trained = False
                return
            
            # 训练每个红球模型
            for i, model in enumerate(self.red_models):
                print(f"   训练红球模型 {i+1}/{len(self.red_models)}")
                
                # 为每个模型添加轻微的特征变化
                varied_features = features + np.random.randn(*features.shape) * 0.01 * (i + 1)
                
                try:
                    model.fit(varied_features, labels)
                except Exception as e:
                    print(f"   [WARNING] 模型 {i+1} 训练失败: {e}")
            
            self.trained = True
            print(f"[INFO] {self.algorithm_name} 训练完成")
            
        except Exception as e:
            print(f"[ERROR] {self.algorithm_name} 训练失败: {e}")
            self.trained = False
    
    def predict_single_set(self, recent_data, pattern_analysis=None):
        """生成单个预测集合"""
        try:
            # 准备预测特征
            features = self.feature_extractor.prepare_prediction_features(recent_data, pattern_analysis)
            
            # 收集所有模型的预测
            all_red_predictions = []
            
            for i, model in enumerate(self.red_models):
                try:
                    if hasattr(model, 'predict_proba') and self.trained:
                        # 获取概率预测
                        probs = model.predict_proba(features.reshape(1, -1))[0]
                        
                        # 转换为号码预测
                        pred_numbers = self._probs_to_numbers(probs, 2, min_val=1, max_val=33)
                        all_red_predictions.extend(pred_numbers)
                        
                    else:
                        # 备用预测
                        pred_numbers = np.random.choice(range(1, 34), 2, replace=False)
                        all_red_predictions.extend(pred_numbers)
                        
                except Exception as e:
                    print(f"   模型 {i+1} 预测失败: {e}")
                    # 添加备用预测
                    backup_numbers = np.random.choice(range(1, 34), 2, replace=False)
                    all_red_predictions.extend(backup_numbers)
            
            # 选择最终红球号码
            final_red = self._select_best_combination(all_red_predictions, 6)
            
            # 蓝球预测（简化实现）
            blue_numbers = [np.random.randint(1, 17)]
            
            return {
                'red': sorted(final_red),
                'blue': blue_numbers,
                'algorithm': self.algorithm_name,
                'confidence': 0.75  # 固定置信度
            }
            
        except Exception as e:
            print(f"[ERROR] {self.algorithm_name} 预测失败: {e}")
            return self._generate_fallback_prediction()
    
    def _probs_to_numbers(self, probabilities, num_select, min_val=1, max_val=33):
        """将概率转换为号码"""
        try:
            if len(probabilities) < max_val - min_val + 1:
                # 概率数组长度不匹配，使用随机选择
                return np.random.choice(range(min_val, max_val + 1), num_select, replace=False).tolist()
            
            # 选择概率最高的号码
            top_indices = np.argsort(probabilities)[-num_select * 2:][::-1]
            selected_numbers = []
            
            for idx in top_indices:
                number = (idx % (max_val - min_val + 1)) + min_val
                if number not in selected_numbers:
                    selected_numbers.append(number)
                    if len(selected_numbers) >= num_select:
                        break
            
            return selected_numbers
            
        except Exception:
            return np.random.choice(range(min_val, max_val + 1), num_select, replace=False).tolist()
    
    def _select_best_combination(self, all_predictions, num_select):
        """从所有预测中选择最佳组合"""
        # 统计每个号码的出现频次
        number_counts = defaultdict(int)
        for num in all_predictions:
            if isinstance(num, (list, tuple)):
                for n in num:
                    if 1 <= n <= 33:
                        number_counts[n] += 1
            elif 1 <= num <= 33:
                number_counts[num] += 1
        
        # 选择出现频次最高的号码
        sorted_numbers = sorted(number_counts.items(), key=lambda x: x[1], reverse=True)
        selected = [num for num, count in sorted_numbers[:num_select]]
        
        # 如果数量不够，随机补充
        while len(selected) < num_select:
            remaining = [i for i in range(1, 34) if i not in selected]
            if remaining:
                selected.append(np.random.choice(remaining))
            else:
                break
        
        return selected[:num_select]
    
    def _generate_fallback_prediction(self):
        """生成备用预测"""
        red = sorted(np.random.choice(range(1, 34), 6, replace=False))
        blue = [np.random.randint(1, 17)]
        return {
            'red': red, 
            'blue': blue,
            'algorithm': self.algorithm_name,
            'confidence': 0.1
        }

class FeatureExtractor:
    """特征提取器"""
    
    def extract_features_and_labels(self, historical_data):
        """从历史数据提取特征和标签"""
        features = []
        labels = []
        
        try:
            window_size = 5
            
            for i in range(len(historical_data) - window_size):
                # 提取窗口特征
                window = historical_data[i:i+window_size]
                feature_vector = self._extract_window_features(window)
                
                # 提取标签（下一期的第一个红球号码）
                next_period = historical_data[i + window_size]
                if isinstance(next_period, (list, tuple)) and len(next_period) > 0:
                    label = next_period[0] - 1  # 转换为0-32索引
                else:
                    label = 0
                
                features.append(feature_vector)
                labels.append(label)
                
            return np.array(features), np.array(labels)
            
        except Exception as e:
            print(f"特征提取错误: {e}")
            return np.array([]), np.array([])
    
    def _extract_window_features(self, window):
        """从窗口数据提取特征"""
        features = []
        
        try:
            all_numbers = []
            
            # 收集窗口内所有号码
            for period in window:
                if isinstance(period, (list, tuple)):
                    all_numbers.extend([n for n in period if isinstance(n, (int, float))])
                elif isinstance(period, (int, float)):
                    all_numbers.append(period)
            
            if not all_numbers:
                return [0] * 20  # 返回默认特征
            
            # 基本统计特征
            features.extend([
                np.mean(all_numbers),                    # 平均值
                np.std(all_numbers) if len(all_numbers) > 1 else 0,  # 标准差
                len(set(all_numbers)),                   # 唯一数字数量
                max(all_numbers),                        # 最大值
                min(all_numbers),                        # 最小值
                np.median(all_numbers),                  # 中位数
                len(all_numbers),                        # 总数字数量
            ])
            
            # 分布特征
            small_zone = sum(1 for n in all_numbers if 1 <= n <= 11)
            medium_zone = sum(1 for n in all_numbers if 12 <= n <= 22)
            large_zone = sum(1 for n in all_numbers if 23 <= n <= 33)
            
            total_valid = small_zone + medium_zone + large_zone
            if total_valid > 0:
                features.extend([
                    small_zone / total_valid,
                    medium_zone / total_valid,  
                    large_zone / total_valid
                ])
            else:
                features.extend([0.33, 0.33, 0.34])
            
            # 奇偶性特征
            odd_count = sum(1 for n in all_numbers if n % 2 == 1)
            even_count = len(all_numbers) - odd_count
            if len(all_numbers) > 0:
                features.extend([
                    odd_count / len(all_numbers),
                    even_count / len(all_numbers)
                ])
            else:
                features.extend([0.5, 0.5])
            
            # 填充到固定长度20
            while len(features) < 20:
                features.append(0.0)
                
            return features[:20]
            
        except Exception:
            return [0.0] * 20
    
    def prepare_prediction_features(self, recent_data, pattern_analysis):
        """准备预测特征"""
        try:
            if isinstance(recent_data, np.ndarray):
                if len(recent_data.shape) > 1:
                    features = recent_data[-1].flatten()[:15]  # 取最后一行，前15个特征
                else:
                    features = recent_data[-15:]  # 取最后15个值
            else:
                features = np.random.randn(15)  # 备用特征
            
            # 添加模式分析特征
            if pattern_analysis and 'consecutive_probabilities' in pattern_analysis:
                pattern_features = pattern_analysis['consecutive_probabilities'][:5]  # 取前5个
            else:
                pattern_features = [0.03] * 5  # 默认均匀概率
            
            # 组合特征
            combined_features = np.concatenate([features, pattern_features])
            
            # 确保长度为20
            if len(combined_features) < 20:
                combined_features = np.pad(combined_features, (0, 20 - len(combined_features)), 'constant')
            
            return combined_features[:20]
            
        except Exception:
            return np.random.randn(20) * 0.1  # 备用特征